from torchtext.models import T5_BASE_GENERATION
from torchtext.prototype.generate import GenerationUtils

eos_idx = 1
# 预训练模型
t5_base = T5_BASE_GENERATION
transform = t5_base.transform()
model = t5_base.get_model()
model.eval()
# 模型的序列生成器
sequence_generator = GenerationUtils(model)
print("模型加载成功！")

# 句子情感分类
task = "sst2 sentence"
labels = {"1": "negative", "2": "positive"}

# 接受用户的输入
while True:
    text = input("请输入句子：")
    if text == "exit":
        break
    input_text = task + ": " + text
    model_input = transform([input_text])
    model_output = sequence_generator.generate(
        model_input, eos_idx=eos_idx, num_beams=1
    )
    output_text = transform.decode(model_output.tolist())
    # 展示预测结果
    print(f"{output_text[0]}\n")
